package org.github.manx98.jlibhac.tools.fs;

import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.HexUtil;
import cn.hutool.crypto.SecureUtil;
import org.github.manx98.jlibhac.common.keys.KeySet;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.fs.IStorage;
import org.github.manx98.jlibhac.fs.SubStorage;

import java.io.IOException;
import java.math.BigInteger;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 1:30
 */
public class XciHeader {
    private final IStorage baseStorage;
    private final KeySet keySet;
    public final byte[] signature;
    public final String magic;

    public final static int SignatureSize = 0x100;
    public final static String HeaderMagic = "HEAD";
    public final int romAreaStartPage;
    public final int rackupAreaStartPage;
    public final byte kekIndex;
    public final byte titleKeyDecIndex;
    public final GameCardSizeInternal gameCardSize;
    public final int cardHeaderVersion;
    public final GameCardAttribute flags;
    public final Uint64 packageId;
    public final Uint64 validDataEndPage;
    public final byte[] aesCbcIv;
    public final long rootPartitionOffset;
    public final long rootPartitionHeaderSize;
    public final String rootPartitionHeaderHash;
    public final byte[] initialDataHash;
    public final int selSec;
    public final int selT1Key;
    public final int selKey;
    public final int limAreaPage;
    public final byte[] imageHash;
    public byte[] initialDataAuthNonce;
    public byte[] initialDataAuthMac;
    public byte[] initialDataAuthData;
    public byte[] initialDataPackageId;
    public byte[] decryptedTitleKey;
    public byte[] initialData;
    public boolean HasInitialData;
    public final static int HeaderMagicValue = 0x44414548; // HEAD

    public final static int EncryptedHeaderSize = 0x70;
    public final static int GcTitleKeyKekIndexMax = 0x10;

    public XciHeader(KeySet keySet, IStorage storage) throws IOException {
        this.baseStorage = storage;
        this.keySet = keySet;
        IStorage[] result = determineXciSubStorages(baseStorage);
        IStorage keyAreaStorage = result[0];
        IStorage bodyStorage = result[1];
        if (keyAreaStorage != null) {
            IStorage r = keyAreaStorage.asStorage();
            HasInitialData = true;
            this.initialDataPackageId = r.readBytes(8);
            r.readBytes(8);
            this.initialDataAuthData = r.readBytes(0x10);
            this.initialDataAuthMac = r.readBytes(0x10);
            this.initialDataAuthNonce = r.readBytes(0xC);
            r.setPosition(0);
            this.initialData = r.readBytes(512);
        }
        IStorage reader = bodyStorage.asStorage();
        this.signature = reader.readBytes(SignatureSize);
        this.magic = reader.readAscii(4);
        if (!HeaderMagic.equals(this.magic)) {
            throw new IllegalArgumentException("Invalid XCI file: Header magic invalid.");
        }
        reader.setPosition(SignatureSize);
        byte[] sigData = reader.readBytes(SignatureSize);
        reader.setPosition(SignatureSize + 4);
        this.romAreaStartPage = reader.readInt();
        this.rackupAreaStartPage = reader.readInt();
        byte keyIndex = (byte) reader.read();
        this.kekIndex = (byte) (keyIndex >> 4);
        this.titleKeyDecIndex = (byte) (keyIndex & 7);
        this.gameCardSize = GameCardSizeInternal.valueOf(reader.read());
        this.cardHeaderVersion = reader.read();
        this.flags = GameCardAttribute.valueOf(reader.read());
        this.packageId = reader.readUint64();
        this.validDataEndPage = reader.readUint64();
        this.aesCbcIv = reader.readBytes(0x10);
        ArrayUtil.reverse(this.aesCbcIv);
        this.rootPartitionOffset = reader.readLong();
        this.rootPartitionHeaderSize = reader.readLong();
        this.rootPartitionHeaderHash = HexUtil.encodeHexStr(reader.readBytes(0x20));
        this.initialDataHash = reader.readBytes(0x20);
        this.selSec = reader.readInt();
        this.selT1Key = reader.readInt();
        this.selKey = reader.readInt();
        this.limAreaPage = reader.readInt();
        this.imageHash = SecureUtil.sha256().digest(sigData);
        reader.setPosition(rootPartitionOffset);
        byte[] headerBytes = reader.readBytes((int) rootPartitionHeaderSize);
        String actualHeaderHash = SecureUtil.sha256().digestHex(headerBytes);
        if (!actualHeaderHash.equals(rootPartitionHeaderHash)) {
            throw new IllegalArgumentException("Invalid XCI, header hash=" + rootPartitionHeaderHash + ", actual=" + actualHeaderHash);
        }
        byte[] key = new byte[0x10];
        if (decryptCardInitialData(key, initialData, kekIndex, keySet)) {
            this.decryptedTitleKey = key;
        }
    }

    public boolean decryptCardInitialData(byte[] dest, byte[] initialData, int kekIndex, KeySet keySet) {
        if (initialData == null || initialData.length == 0) {
            return false;
        }
        if (kekIndex >= GcTitleKeyKekIndexMax) {
            return false;
        }
        if (keySet.GcTitleKeyKeks[kekIndex].isZeros()) {
            return false;
        }
        throw new UnsupportedOperationException("Not supported yet");
    }

    private IStorage[] determineXciSubStorages(IStorage baseStorage) throws IOException {
        long storageSize = baseStorage.getSize();
        if (storageSize >= 0x1104) {
            int magic = baseStorage.readInt();
            if (magic == HeaderMagicValue) {
                IStorage keyAreaStorage = new SubStorage(baseStorage, 0, 0x1000);
                IStorage bodyStorage = new SubStorage(baseStorage, 0x1000, storageSize - 0x1000);
                return new IStorage[]{keyAreaStorage, bodyStorage};
            }
        }
        return new IStorage[]{null, baseStorage.asStorage()};
    }
}
